# -*- coding: utf-8 -*-
# @Time    : 2021/9/3 15:18
# @Author  : huangwei
# @File    : method.py
# @Software: PyCharm
import copy
import math
import os
import traceback

import cv2
import ffmpeg
import numpy as np
from multiprocessing import Pool

from functions import water_psnr

BLOCK_SHAPE = np.array((8, 8))  # 进行分块的大小
INSERT_ROW = BLOCK_SHAPE[0] - 1
INSERT_COL = BLOCK_SHAPE[1] - 1
MARKSTRENGTH = 100  # 值越大 鲁棒性越强，但是对原文件影响越大

PROCESS_NUM = 3  # 最大进程数
GAP = 30  # 视频分段长度
GAP_TIME = 1  # 插入水印的长度

EXTRACT_MAX_LENGTH = 35
EXTRACT_GAP_TIME = 5
EXTRACT_MAX_GAP_NUM = int(EXTRACT_MAX_LENGTH / EXTRACT_GAP_TIME) - 1

# 加入水印的标志，用于判断是否加入了水印
FLAG = np.array(
    [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
     1,
     0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0,
     1, 0, 1, 0, 1])


def find_dup( numbers ):
    """判断数组中是否存在值相等的"""
    numbers = sorted(numbers)
    for i in range(len(numbers) - 1):
        if numbers[i] == numbers[i + 1]:
            return True

    return False


def get_video_param( video_path ):
    """
    返回视频流中的元素
    :param video_path:
    :return: {duration, width, height, codec_name, pix_fmt, fps, audio}
    """
    probe = ffmpeg.probe(video_path)
    duration = float(probe['format']['duration'])

    video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)

    if video_stream is None:
        raise AssertionError("please check your file, this may not be a video!")

    audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)

    if audio_stream is None:
        is_audio_exists = False
    else:
        is_audio_exists = True

    width = video_stream['width']
    height = video_stream['height']
    codec_name = video_stream['codec_name']
    pix_fmt = video_stream['pix_fmt']
    frame_rate = video_stream['avg_frame_rate']
    fps = round(eval(frame_rate), 2)

    return {"duration": duration, "width": width, "height": height, "codec_name": codec_name, "pix_fmt": pix_fmt,
            "fps": fps, "audio": is_audio_exists}


def extract_from_video( input_filepath, output_filepath, watermark_path, tmp_directory ):
    """
    查看视频中是否存在水印
    """
    # 取出视频流的多个参数不用每次都取提高速度
    video_param = get_video_param(input_filepath)
    duration = video_param['duration']

    # 将视频分为 5s 一个段，最后一段 5 到 10s 长
    # 视频最短为 5s，即 gap_num >= 1
    if duration < EXTRACT_MAX_LENGTH:
        gap_num = int(duration / EXTRACT_GAP_TIME)
    else:
        gap_num = EXTRACT_MAX_GAP_NUM

    # 多进程
    process_num = min(gap_num, PROCESS_NUM)
    process_num = max(process_num, 1)
    p = Pool(process_num)
    for i in range(gap_num - 1):
        # 这是 5s 整的部分
        start_time = i * EXTRACT_GAP_TIME
        gap_path = "{0}/gap{1}.mp4".format(tmp_directory, i)

        # 截取视频
        os.system(
            "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3} -loglevel quiet".format(start_time, EXTRACT_GAP_TIME,
                                                                                          input_filepath, gap_path))

        process = p.apply_async(gap_extract_watermark, args=(gap_path, output_filepath, watermark_path, video_param,))

        print("5s部分查询完毕！该部分的返回值为：", process.get())
        if process.get():
            return True

    # 剩余的部分  即 5s 到 10s 的部分
    start_time = (gap_num - 1) * EXTRACT_GAP_TIME
    gap_path = "{0}/gap{1}.mp4".format(tmp_directory, (gap_num - 1))
    # 截取视频
    os.system(
        "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3} -loglevel quiet".format(start_time, EXTRACT_GAP_TIME * 2,
                                                                                      input_filepath, gap_path))
    process = p.apply_async(gap_extract_watermark, args=(gap_path, output_filepath, watermark_path, video_param,))

    print("剩余部分查询完毕！该部分的返回值为：", process.get())
    if process.get():
        return True
    p.close()  # 关闭进程池，表示不能再往进程池中添加进程，需要在join之前调用
    p.join()  # 等待进程池中的所有进程执行完毕

    print("未查找到水印！")
    return False


def gap_extract_watermark( input_filepath, output_filepath, watermark_path, video_param ):
    """
    从小段视频中提取水印
    """
    # 转为 yuv 格式
    (tmp_path, tmp_filename) = os.path.split(input_filepath)  # 分离路径和文件名
    tmp_name, _ = os.path.splitext(tmp_filename)  # 分离文件名和后缀
    yuv_path = "{0}/{1}.yuv".format(tmp_path, tmp_name)

    os.system("ffmpeg -y -vcodec {0} -i {1} -pix_fmt {2} {3} -loglevel quiet".format(video_param['codec_name'],
                                                                                     input_filepath,
                                                                                     video_param['pix_fmt'], yuv_path))
    exist_watermark = find_flag(yuv_path, (video_param['height'], video_param['width']), output_filepath,
                                watermark_path)

    return exist_watermark


def find_flag( input_filepath, img_shape, output_filepath, WATERMARK_PATH ):
    fp = open(input_filepath, "rb")  # 读取yuv文件
    input_data = fp.read()

    file_length = len(input_data)  # 整个文件字符长度
    y_length = img_shape[0] * img_shape[1]  # 单帧的 y 向量大小
    yuv_length = int(y_length * 3 / 2)  # 单帧的 yuv 向量的大小
    frames_num = int(file_length / yuv_length)  # 总帧数

    print("6.开始判断该段yuv中的视频帧是否存在水印。视频帧总数为：", frames_num)

    frame = bytearray(yuv_length)  # 用于存储单帧的 yuv 数据

    # 用于存储单帧的 y 向量数据
    # 创建一个 img_shape[0] 行 img_shape[1] 列的二维数组
    data = [[0 for x in range(img_shape[1])] for y in range(img_shape[0])]

    # 用于将 y 向量划分成 8*8 的 block
    strides = 4 * np.array([img_shape[1] * BLOCK_SHAPE[0], BLOCK_SHAPE[1], img_shape[1], 1])
    frame_block_shape = (img_shape[0] // BLOCK_SHAPE[0], img_shape[1] // BLOCK_SHAPE[1], BLOCK_SHAPE[0], BLOCK_SHAPE[1])

    # 初始化 block 的 index
    block_index = [(i, j) for i in range(frame_block_shape[0]) for j in range(frame_block_shape[1])]

    # 用来保存psnr的最大值
    psnr_max = 0
    # 用来记录有几帧的 FLAG 标记大于13
    mark_num = 0

    # 读取每一帧的数据
    # 如果每一帧都不存在水印则无水印
    # 否则找到一帧 flag>13 且 满足提取的水印长宽和原来的长宽是一致的
    for i in range(frames_num):
        try:
            if i % 10 == 0:
                print("frame:", i)
            for j in range(yuv_length):
                frame[j] = input_data[i * yuv_length + j]

            # 取出该帧的 y 向量数据
            y_index = 0
            for row in range(img_shape[0]):
                for col in range(img_shape[1]):
                    data[row][col] = frame[y_index]
                    y_index += 1

            frame_y = np.array(data)

            # 将 y 向量划分为 8 * 8 的块
            frame_block = np.lib.stride_tricks.as_strided(frame_y.astype(np.float32), frame_block_shape, strides)

            # 提取出前80 个 block 中插入的数据
            temp_list = []
            for k in range(80):
                frame_block_dct = cv2.dct(frame_block[block_index[k]])
                temp_list.append(frame_block_dct[INSERT_ROW][INSERT_COL] % MARKSTRENGTH)
                frame_block[block_index[k]] = cv2.idct(frame_block_dct)

            list_bit = np.array(temp_list) > 15
            # print(list_bit)

            # 将80 转成 16 去除误差
            temp_bit = []
            for k in range(16):
                temp = int(list_bit[k]) + int(list_bit[k + 16]) + int(list_bit[k + 16 * 2]) + int(
                    list_bit[k + 16 * 3]) + int(list_bit[k + 16 * 4])
                temp_bit.append(temp)

            flag_bit = np.array(temp_bit) > 2

            count_start = 0
            for k in range(16):
                if flag_bit[k] == FLAG[k]:
                    count_start += 1

            if count_start > 13:
                # 则说明该帧一定概率上存在水印，提取出来
                # 再提取 81 到 180 个 block 的水印属性信息
                print("find one frame may exist watermark!")
                mark_num += 1
                temp_list = []
                for k in range(80, 180):
                    frame_block_dct = cv2.dct(frame_block[block_index[k]])
                    temp_list.append(frame_block_dct[INSERT_ROW][INSERT_COL] % MARKSTRENGTH)
                    frame_block[block_index[k]] = cv2.idct(frame_block_dct)

                list_bit = np.array(temp_list) > 15

                height_arr = []
                width_arr = []
                for k in range(5):
                    h_tmp = list_bit[0 + k * 20] * 512 + list_bit[1 + k * 20] * 256 + list_bit[2 + k * 20] * 128 + \
                            list_bit[
                                3 + k * 20] * 64 + list_bit[4 + k * 20] * 32 + list_bit[5 + k * 20] * 16 + list_bit[
                                6 + k * 20] * 8 + \
                            list_bit[7 + k * 20] * 4 + list_bit[8 + k * 20] * 2 + list_bit[9 + k * 20]
                    height_arr.append(h_tmp)

                    w_tmp = list_bit[10 + k * 20] * 512 + list_bit[11 + k * 20] * 256 + list_bit[12 + k * 20] * 128 + \
                            list_bit[
                                13 + k * 20] * 64 + list_bit[14 + k * 20] * 32 + list_bit[15 + k * 20] * 16 + list_bit[
                                16 + k * 20] * 8 + list_bit[17 + k * 20] * 4 + list_bit[18 + k * 20] * 2 + list_bit[
                                19 + k * 20]
                    width_arr.append(w_tmp)

                # 宽高中必有值相等的数，否则当作无水印，下一帧
                print("width arr:", width_arr)
                print("height arr:", height_arr)
                if find_dup(height_arr) and find_dup(width_arr):
                    height = np.argmax(np.bincount(height_arr))
                    width = np.argmax(np.bincount(width_arr))
                    print("watermark size:", width, height)

                    watermark_size = width * height
                    insert_size = watermark_size + len(FLAG) + 100

                    if insert_size > len(block_index):
                        print("水印标志存在但无法提取出完整的水印信息！所以无法提取出水印！")
                        return True

                    # 提取水印信息
                    temp_list = []
                    for k in range(180, 180 + watermark_size):
                        frame_block_dct = cv2.dct(frame_block[block_index[k]])
                        temp_list.append(frame_block_dct[INSERT_ROW][INSERT_COL] % MARKSTRENGTH)
                        frame_block[block_index[k]] = cv2.idct(frame_block_dct)

                    list_bit = np.array(temp_list) > 15

                    watermark_data = [[0 for x in range(width)] for y in range(height)]

                    index = 0
                    for row in range(height):
                        for col in range(width):
                            watermark_data[row][col] = int(list_bit[index])
                            index += 1

                    watermark = 255 * np.array(watermark_data)

                    psnr_value = water_psnr(WATERMARK_PATH, watermark)
                    # 如果psnr > 15 直接保存返回
                    # 否则找出最大的 psnr 时保存返回
                    if psnr_value > 15:
                        print("视频中存在水印，将水印提取保存在 %s" % output_filepath)
                        cv2.imwrite(output_filepath, watermark)
                        return True
                    elif psnr_value > psnr_max:
                        psnr_max = psnr_value
                        cv2.imwrite(output_filepath, watermark)  # 覆盖之前的
        except Exception as e:
            print("there find a exception:", e)
            print(traceback.print_exc())
            print("next frame!")

    if psnr_max > 0:
        print("视频中存在水印，但psnr较小为：{0}，将水印提取保存在 {1}".format(psnr_max, output_filepath))
        return True

    if mark_num > 5:
        print("该视频中标记帧数大于5，有可能存在水印但无法成功提取出来。或提取水印算法不适合该视频数据,请换一段视频。将往该视频插入水印。")
        return False

    return False


def split_video( input_filepath, output_filepath, watermark_filepath, video_param, tmp_path ):
    """
    将视频进行分段，每一段开启一个进程
    :param output_filepath:
    :param watermark_filepath:
    :param video_param:
    :param input_filepath:
    :param tmp_path:
    :return:
    """

    # 视频时长过小情况不考虑，即 GAP > 3
    # 最后一段的长度为 3s 到 gap+3s 之间
    duration = video_param['duration']
    if duration % GAP < 1.5:
        gap_num = int(duration / GAP) - 1
    else:
        gap_num = int(duration / GAP)

    # 长为 30s 的视频段
    # 对Pool对象调用join() 方法会等待所有子进程执行完毕，调用join()之前必须先调用close()，让其不再接受新的Process了
    process_num = min((gap_num + 1), PROCESS_NUM)
    process_num = max(process_num, 1)
    p = Pool(process_num)

    for i in range(gap_num):
        start_time = i * GAP
        # 插入水印和无变化的两部分
        gap_path1 = "{0}/gap{1}_1.mp4".format(tmp_path, i)
        gap_path2 = "{0}/gap{1}_2.mp4".format(tmp_path, i)

        # 插完水印后的文件路径
        gap_insert_path1 = "{0}/gap{1}_insert.mp4".format(tmp_path, i)

        # 分别截取两部分
        os.system(
            "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3} -loglevel quiet".format(start_time, GAP_TIME,
                                                                                          input_filepath, gap_path1))
        os.system(
            "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3} -loglevel quiet".format(start_time + GAP_TIME,
                                                                                          GAP - GAP_TIME,
                                                                                          input_filepath, gap_path2))

        # 往第一部分插入水印
        # gap_insert_watermark(gap_path1, gap_insert_path1, watermark_filepath, video_param)
        process = p.apply_async(gap_insert_watermark,
                                args=(gap_path1, gap_insert_path1, watermark_filepath, video_param,))

        print("30s视频段划分完成并插入水印！")

    # 不够 一个 GAP 的部分
    start_time = gap_num * GAP
    gap_path1 = "{0}/gap{1}_1.mp4".format(tmp_path, gap_num)
    gap_path2 = "{0}/gap{1}_2.mp4".format(tmp_path, gap_num)

    gap_insert_path1 = "{0}/gap{1}_insert.mp4".format(tmp_path, gap_num)
    # 分别截取两部分
    os.system(
        "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3} -loglevel quiet".format(start_time, GAP_TIME,
                                                                                      input_filepath, gap_path1))
    os.system(
        "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3} -loglevel quiet".format(start_time + GAP_TIME,
                                                                                      GAP,
                                                                                      input_filepath, gap_path2))
    # gap_insert_watermark(gap_path1, gap_insert_path1, watermark_filepath, video_param)
    process = p.apply_async(gap_insert_watermark,
                            args=(gap_path1, gap_insert_path1, watermark_filepath, video_param,))
    p.close()
    p.join()
    print("不够30s的部分划分和插入水印完成！")

    # 写一个 txt 文件用于连接视频
    list_path = "{0}/list.txt".format(tmp_path)
    with open(list_path, 'a+') as fp:
        for i in range(gap_num + 1):
            fp.write("file gap{0}_insert.mp4\n".format(i))
            fp.write("file gap{0}_2.mp4\n".format(i))

    # 视频合并
    os.system("ffmpeg -y -f concat -safe 0 -i {0} -c copy {1}  -loglevel quiet".format(list_path, output_filepath))
    print("视频合并完成：", output_filepath)


def gap_insert_watermark( input_filepath, output_filepath, watermark_filepath, video_param ):
    """
    往该段时间所有帧插入水印
    1. 将视频转换为 yuv 格式
    2. 对yuv插入水印
    3. yuv 转 mp4
    :param watermark_filepath:
    :param video_param:
    :param input_filepath:
    :param output_filepath:
    :return:
    """

    (tmp_path, tmp_filename) = os.path.split(input_filepath)  # 分离路径和文件名
    tmp_name, _ = os.path.splitext(tmp_filename)  # 分离文件名和后缀
    yuv_path = "{0}/{1}.yuv".format(tmp_path, tmp_name)
    yuv_embed_path = "{0}/{1}_embed.yuv".format(tmp_path, tmp_name)

    os.system("ffmpeg -y -vcodec {0} -i {1} -pix_fmt {2} {3} -loglevel quiet".format(video_param['codec_name'],
                                                                                     input_filepath,
                                                                                     video_param['pix_fmt'], yuv_path))

    insert2yuv(yuv_path, yuv_embed_path, watermark_filepath, (video_param['height'], video_param['width']))

    os.system(
        "ffmpeg -y -s:v {0}x{1} -pix_fmt {2} -r {3} -i {4} {5} -loglevel quiet".format(video_param['width'],
                                                                                       video_param['height'],
                                                                                       video_param['pix_fmt'],
                                                                                       video_param['fps'],
                                                                                       yuv_embed_path,
                                                                                       output_filepath))


def insert2yuv( input_filepath, output_filepath, watermark_filepath, img_shape ):
    fp = open(input_filepath, "rb")  # 读取yuv文件
    input_data = fp.read()

    file_length = len(input_data)  # 整个文件字符长度
    y_length = img_shape[0] * img_shape[1]  # 单帧的 y 向量大小
    yuv_length = int(y_length * 3 / 2)  # 单帧的 yuv 向量的大小
    frames_num = int(file_length / yuv_length)  # 总帧数
    print("需要插入的视频帧总数为：", frames_num / 2)

    frame = bytearray(yuv_length)  # 用于存储单帧的 yuv 数据

    # 用于存储单帧的 y 向量数据
    # 创建一个 img_shape[0] 行 img_shape[1] 列的二维数组
    data = [[0 for x in range(img_shape[1])] for y in range(img_shape[0])]

    # 用于将 y 向量划分成 8*8 的 block
    strides = 4 * np.array([img_shape[1] * BLOCK_SHAPE[0], BLOCK_SHAPE[1], img_shape[1], 1])
    frame_block_shape = (img_shape[0] // BLOCK_SHAPE[0], img_shape[1] // BLOCK_SHAPE[1], BLOCK_SHAPE[0], BLOCK_SHAPE[1])

    # 初始化 block 的 index
    block_index = [(i, j) for i in range(frame_block_shape[0]) for j in range(frame_block_shape[1])]

    # 生成水印 bit 信息和属性 bit 信息，即将要插入的数据转为二进制
    watermark = cv2.imread(watermark_filepath, cv2.IMREAD_GRAYSCALE)

    watermark_size = watermark.shape[0] * watermark.shape[1]
    insert_size = watermark_size + len(FLAG) + 100

    # 如果水印过大，则将其进行缩放到合适的大小
    if insert_size > len(block_index):
        print("最多可嵌入{}kb信息，水印大小{}kb，因此将水印进行缩放".format(len(block_index) / 1024, watermark_size / 1024))
        max_watermark_size = len(block_index) - len(FLAG) - 100
        scale = (max_watermark_size / watermark_size) ** 0.5
        shape0 = int(watermark.shape[0] * scale)
        shape1 = int(watermark.shape[1] * scale)
        watermark_size = shape0 * shape1

        watermark = cv2.resize(watermark, (shape1, shape0))
        print("新的水印大小为{0}*{1}。".format(watermark.shape[0], watermark.shape[1]))

    watermark_bit = watermark.flatten() > 128

    height = format(watermark.shape[0], "b").zfill(10)  # 将 height, width 转为长度为 10 的二进制数据
    width = format(watermark.shape[1], "b").zfill(10)
    watermark_attr_bit = (height + width) * 5

    # 依次取出每一帧插入水印
    with open(output_filepath, "ab+") as wp:
        for i in range(frames_num):
            for j in range(yuv_length):
                frame[j] = input_data[i * yuv_length + j]

            if i % 2 == 0:
                # 取出该帧的 y 向量数据
                y_index = 0
                for row in range(img_shape[0]):
                    for col in range(img_shape[1]):
                        data[row][col] = frame[y_index]
                        y_index += 1

                frame_y = np.array(data)

                # 将 y 向量划分为 8 * 8 的块
                frame_block = np.lib.stride_tricks.as_strided(frame_y.astype(np.float32), frame_block_shape, strides)
                embed_frame = copy.deepcopy(frame_y)

                print("数据帧插入：", i)

                # 对前 80 个块进行标志位嵌入
                for k in range(len(FLAG)):
                    # 对 frame_block[k] 进行离散余弦变换获得DCT系数矩阵
                    # 运用余数定理实现水印嵌入
                    frame_block_dct = cv2.dct(frame_block[block_index[k]])
                    # print("frame_block_dct:", frame_block_dct[INSERT_ROW][INSERT_COL])

                    if FLAG[k] == 0:
                        frame_block_dct[INSERT_ROW][INSERT_COL] = math.floor(
                            frame_block_dct[INSERT_ROW][INSERT_COL] / MARKSTRENGTH) * MARKSTRENGTH + 10
                    else:
                        frame_block_dct[INSERT_ROW][INSERT_COL] = math.floor(
                            frame_block_dct[INSERT_ROW][INSERT_COL] / MARKSTRENGTH) * MARKSTRENGTH + 30

                        # 先将值的范围限定到 0 到 255， 再四舍五入取整
                        frame_block[block_index[k]] = np.rint(np.clip(cv2.idct(frame_block_dct), a_min=0, a_max=255))

                # 嵌入水印属性信息 100 bits
                for k in range(len(watermark_attr_bit)):

                    frame_block_dct = cv2.dct(frame_block[block_index[k + len(FLAG)]])

                    if int(watermark_attr_bit[k]):
                        frame_block_dct[INSERT_ROW][INSERT_COL] = math.floor(
                            frame_block_dct[INSERT_ROW][INSERT_COL] / MARKSTRENGTH) * MARKSTRENGTH + 30
                    else:
                        frame_block_dct[INSERT_ROW][INSERT_COL] = math.floor(
                            frame_block_dct[INSERT_ROW][INSERT_COL] / MARKSTRENGTH) * MARKSTRENGTH + 10

                    frame_block[block_index[k + len(FLAG)]] = np.rint(
                        np.clip(cv2.idct(frame_block_dct), a_min=0, a_max=255))

                # 嵌入水印信息
                for k in range(watermark_size):

                    frame_block_dct = cv2.dct(frame_block[block_index[k + len(FLAG) + len(watermark_attr_bit)]])

                    if watermark_bit[k] == 0:
                        frame_block_dct[INSERT_ROW][INSERT_COL] = math.floor(
                            frame_block_dct[INSERT_ROW][INSERT_COL] / MARKSTRENGTH) * MARKSTRENGTH + 10
                    else:
                        frame_block_dct[INSERT_ROW][INSERT_COL] = math.floor(
                            frame_block_dct[INSERT_ROW][INSERT_COL] / MARKSTRENGTH) * MARKSTRENGTH + 30

                    frame_block[block_index[k + len(FLAG) + len(watermark_attr_bit)]] = np.rint(
                        np.clip(cv2.idct(frame_block_dct), a_min=0, a_max=255))

                # 四维转为二维，还少了整除剩下的部分。
                part_frame = np.concatenate(np.concatenate(frame_block, 1), 1)

                # 将插入水印的部分放回原处补齐整除剩下的部分
                embed_frame[:part_frame.shape[0], :part_frame.shape[1]] = part_frame

                # 写回 yuv
                y_index = 0
                for row in range(img_shape[0]):
                    for col in range(img_shape[1]):
                        frame[y_index] = int(embed_frame[row][col])
                        y_index += 1

            wp.write(frame)
